from dataclasses import dataclass, field
from typing import List, Optional

import numpy as np
import pandas as pd
from . import base


@dataclass
class DropNA(base.BaseTransformer, base.CopyMixin):
    cols: Optional[List[str]] = field(default_factory=lambda: ['close'])

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data = data.dropna(subset=self.cols) if self.cols else data.dropna(subset=None)
        return data


@dataclass
class FillNA(base.BaseTransformer, base.CopyMixin):
    cols: List[str] = field(default_factory=lambda: ['close'])
    method: str = 'ffill'
    fill_args: dict = field(default_factory=dict)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data.loc[:, self.cols] = data.loc[:, self.cols].fillna(method=self.method, **self.fill_args)  # type: ignore
        return data


@dataclass
class Ffill(base.BaseTransformer, base.CopyMixin):
    cols: List[str] = field(default_factory=lambda: ['close'])

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data.loc[:, self.cols] = data.loc[:, self.cols].ffill()  # type: ignore
        return data


@dataclass
class FillInf(base.BaseTransformer, base.CopyMixin):
    cols: List[str] = field(default_factory=lambda: ['close'])
    positive: bool = True
    method: Optional[str] = 'ffill'
    fill_args: dict = field(default_factory=dict)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        if self.positive:
            cols = data.columns[(data == float('inf')).any(axis=0)]
        else:
            cols = data.columns[(data == -float('inf')).any(axis=0)]
        for col in cols:
            if data[col].isna().any():
                raise ValueError(f'{col} contains NA')
            if self.positive:
                data.loc[data[col] == float('inf'), col] = np.nan
            else:
                data.loc[data[col] == -float('inf'), col] = np.nan
            data[col] = data[col].fillna(method=self.method, **self.fill_args)
        return data
